{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xwpZkPfDBoO-"
      },
      "source": [
        "# WordPiece tokenization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zbtIWnqcBoPE"
      },
      "source": [
        "Install the Transformers, Datasets, and Evaluate libraries to run this notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "关于wordpiece算法原理的详细文档描述：https://huggingface.co/learn/nlp-course/en/chapter6/6?fw=pt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "uGJHoXw1BoPF"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: datasets in e:\\programs\\anaconda\\lib\\site-packages (2.21.0)\n",
            "Requirement already satisfied: evaluate in e:\\programs\\anaconda\\lib\\site-packages (0.4.2)\n",
            "Requirement already satisfied: transformers[sentencepiece] in e:\\programs\\anaconda\\lib\\site-packages (4.44.0)\n",
            "Requirement already satisfied: filelock in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (3.13.1)\n",
            "Requirement already satisfied: numpy>=1.17 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (1.26.4)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (17.0.0)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (2.1.4)\n",
            "Requirement already satisfied: requests>=2.32.2 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (4.66.5)\n",
            "Requirement already satisfied: xxhash in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (3.5.0)\n",
            "Requirement already satisfied: multiprocess in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec<=2024.6.1,>=2023.1.0 in e:\\programs\\anaconda\\lib\\site-packages (from fsspec[http]<=2024.6.1,>=2023.1.0->datasets) (2023.10.0)\n",
            "Requirement already satisfied: aiohttp in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (3.9.3)\n",
            "Requirement already satisfied: huggingface-hub>=0.21.2 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (0.24.5)\n",
            "Requirement already satisfied: packaging in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (23.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in e:\\programs\\anaconda\\lib\\site-packages (from datasets) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in e:\\programs\\anaconda\\lib\\site-packages (from transformers[sentencepiece]) (2023.10.3)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in e:\\programs\\anaconda\\lib\\site-packages (from transformers[sentencepiece]) (0.4.4)\n",
            "Requirement already satisfied: tokenizers<0.20,>=0.19 in e:\\programs\\anaconda\\lib\\site-packages (from transformers[sentencepiece]) (0.19.1)\n",
            "Requirement already satisfied: protobuf in e:\\programs\\anaconda\\lib\\site-packages (from transformers[sentencepiece]) (3.20.3)\n",
            "Requirement already satisfied: sentencepiece!=0.1.92,>=0.1.91 in e:\\programs\\anaconda\\lib\\site-packages (from transformers[sentencepiece]) (0.2.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in e:\\programs\\anaconda\\lib\\site-packages (from aiohttp->datasets) (1.2.0)\n",
            "Requirement already satisfied: attrs>=17.3.0 in e:\\programs\\anaconda\\lib\\site-packages (from aiohttp->datasets) (23.1.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in e:\\programs\\anaconda\\lib\\site-packages (from aiohttp->datasets) (1.4.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in e:\\programs\\anaconda\\lib\\site-packages (from aiohttp->datasets) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in e:\\programs\\anaconda\\lib\\site-packages (from aiohttp->datasets) (1.9.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in e:\\programs\\anaconda\\lib\\site-packages (from huggingface-hub>=0.21.2->datasets) (4.9.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in e:\\programs\\anaconda\\lib\\site-packages (from requests>=2.32.2->datasets) (2.0.4)\n",
            "Requirement already satisfied: idna<4,>=2.5 in e:\\programs\\anaconda\\lib\\site-packages (from requests>=2.32.2->datasets) (3.4)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in e:\\programs\\anaconda\\lib\\site-packages (from requests>=2.32.2->datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in e:\\programs\\anaconda\\lib\\site-packages (from requests>=2.32.2->datasets) (2024.2.2)\n",
            "Requirement already satisfied: colorama in e:\\programs\\anaconda\\lib\\site-packages (from tqdm>=4.66.3->datasets) (0.4.6)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in e:\\programs\\anaconda\\lib\\site-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in e:\\programs\\anaconda\\lib\\site-packages (from pandas->datasets) (2023.3.post1)\n",
            "Requirement already satisfied: tzdata>=2022.1 in e:\\programs\\anaconda\\lib\\site-packages (from pandas->datasets) (2023.3)\n",
            "Requirement already satisfied: six>=1.5 in e:\\programs\\anaconda\\lib\\site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install datasets evaluate transformers[sentencepiece]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "nkNP9fNiBoPH"
      },
      "outputs": [],
      "source": [
        "corpus = [\n",
        "    \"This is the Hugging Face Course.\",\n",
        "    \"This chapter is about tokenization.\",\n",
        "    \"This section shows several tokenizer algorithms.\",\n",
        "    \"Hopefully, you will be able to understand how they are trained and generate tokens.\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "nitqpdYVBoPI"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "e:\\Programs\\Anaconda\\Lib\\site-packages\\transformers\\tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "OLDzKkNZBoPJ",
        "outputId": "7fc98df2-b1aa-4114-c133-ce409adf01c8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "defaultdict(int,\n",
              "            {'This': 3,\n",
              "             'is': 2,\n",
              "             'the': 1,\n",
              "             'Hugging': 1,\n",
              "             'Face': 1,\n",
              "             'Course': 1,\n",
              "             '.': 4,\n",
              "             'chapter': 1,\n",
              "             'about': 1,\n",
              "             'tokenization': 1,\n",
              "             'section': 1,\n",
              "             'shows': 1,\n",
              "             'several': 1,\n",
              "             'tokenizer': 1,\n",
              "             'algorithms': 1,\n",
              "             'Hopefully': 1,\n",
              "             ',': 1,\n",
              "             'you': 1,\n",
              "             'will': 1,\n",
              "             'be': 1,\n",
              "             'able': 1,\n",
              "             'to': 1,\n",
              "             'understand': 1,\n",
              "             'how': 1,\n",
              "             'they': 1,\n",
              "             'are': 1,\n",
              "             'trained': 1,\n",
              "             'and': 1,\n",
              "             'generate': 1,\n",
              "             'tokens': 1})"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from collections import defaultdict\n",
        "\n",
        "word_freqs = defaultdict(int)\n",
        "for text in corpus:\n",
        "    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)\n",
        "    new_words = [word for word, offset in words_with_offsets]\n",
        "    for word in new_words:\n",
        "        word_freqs[word] += 1\n",
        "\n",
        "word_freqs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ckL961h2BoPK",
        "outputId": "8f16b102-1f3b-4fa6-8af2-bc78d673c0ab"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['##a', '##b', '##c', '##d', '##e', '##f', '##g', '##h', '##i', '##k', '##l', '##m', '##n', '##o', '##p', '##r', '##s', '##t', '##u', '##v', '##w', '##y', '##z', ',', '.', 'C', 'F', 'H', 'T', 'a', 'b', 'c', 'g', 'h', 'i', 's', 't', 'u', 'w', 'y']\n"
          ]
        }
      ],
      "source": [
        "# 首字母作为初始词汇, ##，表示不是词首\n",
        "alphabet = []\n",
        "for word in word_freqs.keys():\n",
        "    if word[0] not in alphabet:\n",
        "        alphabet.append(word[0])\n",
        "    for letter in word[1:]:\n",
        "        if f\"##{letter}\" not in alphabet:\n",
        "            alphabet.append(f\"##{letter}\")\n",
        "\n",
        "alphabet.sort()\n",
        "alphabet\n",
        "\n",
        "print(alphabet)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MbblglExBoPL"
      },
      "outputs": [],
      "source": [
        "# 特殊词元token\n",
        "vocab = [\"[PAD]\", \"[UNK]\", \"[CLS]\", \"[SEP]\", \"[MASK]\"] + alphabet.copy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fICmOiWRBoPL"
      },
      "outputs": [],
      "source": [
        "# 计算splits\n",
        "splits = {\n",
        "    word: [c if i == 0 else f\"##{c}\" for i, c in enumerate(word)]\n",
        "    for word in word_freqs.keys()\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e2GGIqukBoPM"
      },
      "outputs": [],
      "source": [
        "#计算token的频率和相邻token pair的频率，并最后算pair的得分\n",
        "def compute_pair_scores(splits):\n",
        "    letter_freqs = defaultdict(int)\n",
        "    pair_freqs = defaultdict(int)\n",
        "    for word, freq in word_freqs.items():\n",
        "        split = splits[word]\n",
        "        if len(split) == 1:\n",
        "            letter_freqs[split[0]] += freq\n",
        "            continue\n",
        "        for i in range(len(split) - 1):\n",
        "            pair = (split[i], split[i + 1])\n",
        "            letter_freqs[split[i]] += freq\n",
        "            pair_freqs[pair] += freq\n",
        "        letter_freqs[split[-1]] += freq\n",
        "\n",
        "    scores = {\n",
        "        pair: freq / (letter_freqs[pair[0]] * letter_freqs[pair[1]])\n",
        "        for pair, freq in pair_freqs.items()\n",
        "    }\n",
        "    return scores"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "Kqq6XicYBoPN",
        "outputId": "a2fef3b6-e89a-45d5-d784-f2271d48fdeb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('T', '##h'): 0.125\n",
            "('##h', '##i'): 0.03409090909090909\n",
            "('##i', '##s'): 0.02727272727272727\n",
            "('i', '##s'): 0.1\n",
            "('t', '##h'): 0.03571428571428571\n",
            "('##h', '##e'): 0.011904761904761904\n"
          ]
        }
      ],
      "source": [
        "pair_scores = compute_pair_scores(splits)\n",
        "for i, key in enumerate(pair_scores.keys()):\n",
        "    print(f\"{key}: {pair_scores[key]}\")\n",
        "    if i >= 5:\n",
        "        break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A8X_QnL8BoPN",
        "outputId": "5ae5a8eb-e429-49f4-dc0f-629eba2dd69e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('a', '##b') 0.2\n"
          ]
        }
      ],
      "source": [
        "# 计算最高的得分的pair\n",
        "best_pair = \"\"\n",
        "max_score = None\n",
        "for pair, score in pair_scores.items():\n",
        "    if max_score is None or max_score < score:\n",
        "        best_pair = pair\n",
        "        max_score = score\n",
        "\n",
        "print(best_pair, max_score)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "Go-4mWScBoPN"
      },
      "outputs": [],
      "source": [
        "vocab.append(\"ab\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "JDZJWmaTBoPO"
      },
      "outputs": [],
      "source": [
        "def merge_pair(a, b, splits):\n",
        "    for word in word_freqs:\n",
        "        split = splits[word]\n",
        "        if len(split) == 1:\n",
        "            continue\n",
        "        i = 0\n",
        "        while i < len(split) - 1:\n",
        "            if split[i] == a and split[i + 1] == b:\n",
        "                merge = a + b[2:] if b.startswith(\"##\") else a + b\n",
        "                split = split[:i] + [merge] + split[i + 2 :]\n",
        "            else:\n",
        "                i += 1\n",
        "        splits[word] = split\n",
        "    return splits"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "Ht8zveMJBoPO",
        "outputId": "42c2f3db-ada3-44af-c0e2-98d0f6c99d85"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['ab', '##o', '##u', '##t']"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "splits = merge_pair(\"a\", \"##b\", splits)\n",
        "splits[\"about\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_guMdGlTBoPO"
      },
      "outputs": [],
      "source": [
        "#合并最高的相邻的现有词汇，并更新到词汇表里，不断持续这个过程支持达到词汇表定义的长度\n",
        "vocab_size = 70\n",
        "while len(vocab) < vocab_size:\n",
        "    scores = compute_pair_scores(splits)\n",
        "    best_pair, max_score = \"\", None\n",
        "    for pair, score in scores.items():\n",
        "        if max_score is None or max_score < score:\n",
        "            best_pair = pair\n",
        "            max_score = score\n",
        "    splits = merge_pair(*best_pair, splits)\n",
        "    new_token = (\n",
        "        best_pair[0] + best_pair[1][2:]\n",
        "        if best_pair[1].startswith(\"##\")\n",
        "        else best_pair[0] + best_pair[1]\n",
        "    )\n",
        "    vocab.append(new_token)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "Lb966oqKBoPO",
        "outputId": "15325495-4a34-44d5-e8aa-6b691fd68b95"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]', '##a', '##b', '##c', '##d', '##e', '##f', '##g', '##h', '##i', '##k', '##l', '##m', '##n', '##o', '##p', '##r', '##s', '##t', '##u', '##v', '##w', '##y', '##z', ',', '.', 'C', 'F', 'H', 'T', 'a', 'b', 'c', 'g', 'h', 'i', 's', 't', 'u', 'w', 'y', 'ab', 'ab', '##fu', 'Fa', 'Fac', '##ct', '##ful', '##full', '##fully', 'Th', 'ch', '##hm', 'cha', 'chap', 'chapt', '##thm', 'Hu', 'Hug', 'Hugg', 'sh', 'th', 'is', '##thms', '##za', '##zat']\n"
          ]
        }
      ],
      "source": [
        "print(vocab)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BohdOIvpBoPP"
      },
      "outputs": [],
      "source": [
        "#对自然单词进行子词分词\n",
        "def encode_word(word):\n",
        "    tokens = []\n",
        "    while len(word) > 0:\n",
        "        i = len(word)\n",
        "        while i > 0 and word[:i] not in vocab:\n",
        "            i -= 1\n",
        "        if i == 0:\n",
        "            return [\"[UNK]\"]\n",
        "        tokens.append(word[:i])\n",
        "        word = word[i:]\n",
        "        if len(word) > 0:\n",
        "            word = f\"##{word}\"\n",
        "    return tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "qXSltKqSBoPP",
        "outputId": "00093252-a0e9-4a37-8f9d-ccc93f0798a2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['Hugg', '##i', '##n', '##g']\n",
            "['[UNK]']\n"
          ]
        }
      ],
      "source": [
        "print(encode_word(\"Hugging\"))\n",
        "print(encode_word(\"HOgging\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HDklWMHpBoPP"
      },
      "outputs": [],
      "source": [
        "# 对文本进行分词的结果\n",
        "def tokenize(text):\n",
        "    pre_tokenize_result = tokenizer._tokenizer.pre_tokenizer.pre_tokenize_str(text)\n",
        "    pre_tokenized_text = [word for word, offset in pre_tokenize_result]\n",
        "    encoded_words = [encode_word(word) for word in pre_tokenized_text]\n",
        "    return sum(encoded_words, [])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3VMQz1M4BoPP",
        "outputId": "53901981-be5e-443d-9011-9e2f884f8185"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['Th', '##i', '##s', 'is', 'th', '##e', 'Hugg', '##i', '##n', '##g', 'Fac', '##e', 'c', '##o', '##u', '##r', '##s',\n",
              " '##e', '[UNK]']"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tokenize(\"This is the Hugging Face course!\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "WordPiece tokenization",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
